
#include "LPC/lpcsvc.h"
#include "LPC/RealAlpc.h"
#include "LPC/request.h"
#include "watermark/watermark.h"


#include "SafeBox.h"
#include "HookNtdll.h"
#include "HookKernel.h"
#include "HookUser32.h"
#include "HookShell32.h"
#include "HookSechost.h"
#include "header_def.h"

#include <stdio.h>

#include <map>
#include <set>
#include <vector>


SAFEBOX_INFO gBoxInfo = { 0 };
HOST_PROCESS gHostProcess = { 0 };

static BOOL bTest = FALSE;
static BOOL RequestBoxInfo(DWORD dwProcessId);
static BOOL IsRunningAsAdmin();
void Init_SkipFilter();

void ProcessAttach(HMODULE hModule)
{
    char* buffer = new char[MAX_PATH * 3];
    CHAR* strName = buffer + MAX_PATH * 0;
    CHAR* strDir = buffer + MAX_PATH * 1;
    CHAR* strPath = buffer + MAX_PATH * 2;
    
    GetModuleBaseNameA(GetCurrentProcess(), NULL, strName, MAX_PATH);

    if (0 == _stricmp(strName, "safe111.exe") ||
        0 == _stricmp(strName, "svctest.exe"))
    {
        while (FALSE == IsDebuggerPresent())
        {
            Sleep(1000 * 3);
        }
    }

    DWORD dwProcessId = GetCurrentProcessId();;
    DWORD dwThreadId = 0;

    GetModuleFileNameA(hModule, strDir, MAX_PATH);
    LPSTR(strrchr(strDir, '\\'))[1] = '\0';
    SYSTEMTIME systime = { 0 };
    GetLocalTime(&systime);

    _snprintf_s(strPath, MAX_PATH, MAX_PATH, "%slog\\%s_%04d%02d%02d_%02d%02d%02d_%05d.log",
        strDir, strName, systime.wYear, systime.wMonth, systime.wDay,
        systime.wHour, systime.wMinute, systime.wSecond, dwProcessId);

    OPENLOG(strPath);

    Init_SkipFilter();

    Hook_Ntdll();
    Hook_KernelBase();
    Hook_win32u();
    Hook_User32(); // empty
    Hook_Shell32();
    Hook_Sechost();
    Hook_Advapi32();
    Hook_Rpcrt4();// empty

    RequestBoxInfo(dwProcessId);
    delete buffer;
}

void ProcessDetach(HMODULE hModule)
{
    UnHook_Ntdll();
    UnHook_KernelBase();
    UnHook_win32u();
    UnHook_User32();
    UnHook_Shell32();
    UnHook_Sechost();
    UnHook_Advapi32();
    UnHook_Rpcrt4();
    CLOSELOG();
}

static DWORD WINAPI SafeBoxWorkerThread(HANDLE hNotifyEvent)
{
    CLpcSrv Server;
    Server.InitEnv();
    Server.Dispatch(hNotifyEvent);
    Server.UninitEnv();
    return 0;
}

static VOID NTAPI WaitForWorkerThreadStart(_In_ HANDLE hEvent)
{
    if (hEvent)
    {
        WaitForSingleObject(hEvent, -1);
        CloseHandle(hEvent);
    }
    LPC_ProcessStart();
}

static BOOL RequestBoxInfo(DWORD dwProcessId)
{
    BOOLEAN bRet = FALSE;
    BOXINFO_MESSSAGE_REQ Request = { 0 };
    BOXINFO_MESSSAGE_RPL Response = { 0 };

    DWORD dwRecvSize = sizeof(BOXINFO_MESSSAGE_RPL);
    Request.Info.u1.s1.DataLength = sizeof(BOXINFO_MESSSAGE_REQ) - sizeof(LPC_MESSAGE);
    Request.Info.u1.s1.TotalLength = sizeof(BOXINFO_MESSSAGE_REQ);
    Request.dwType = BOXINFO_REQUEST_DATA;
    Request.dwProcessId = dwProcessId;

    if (FALSE == REAL_AlpcClientSendReceive(BOXINFO_RPCNAME, (PLPC_MESSAGE)&Request,
        (PLPC_MESSAGE)&Response, &dwRecvSize, -1))
    {
        return FALSE;
    }

    if (Response.BoxInfo.dwDoaminPid && Response.BoxInfo.dwDoaminPid != dwProcessId)
    {
        return SetSafeBoxInfo(&Response.BoxInfo);
    }
    return TRUE;
}

DWORD WINAPI SetSafeBoxInfo(PSAFEBOX_INFO pBoxInfo)
{
    if (pBoxInfo)
    {
        gBoxInfo = *pBoxInfo;
        if (gBoxInfo.uLpcPort)
        {
            HANDLE hEvent = ::CreateEvent(NULL, FALSE, FALSE, NULL);
            HANDLE hThread = ::CreateThread(NULL, 0, 
                (LPTHREAD_START_ROUTINE)SafeBoxWorkerThread,
                hEvent, 0, NULL);
            CloseHandle(hThread);

            HANDLE hMainThread = GetCurrentThread();
            QueueUserAPC((PAPCFUNC)WaitForWorkerThreadStart,
                hMainThread, (ULONG_PTR)hEvent);
        }
        return sizeof(SAFEBOX_INFO);
    }
    return 0;
}

LPCSTR GetSafeBoxNameA()
{
    return gBoxInfo.strSuffixA;
}

LPCWSTR GetSafeBoxNameW()
{
    return gBoxInfo.strSuffixW;
}

